# modify from https://github.com/mit-han-lab/bevfusion
import torch
import torch.nn as nn
import torch.nn.functional as F
from mmcv.cnn import ConvModule
from mmengine.model import BaseModule

from mmdet3d.registry import MODELS

SWIN_CFG = dict(
    type='mmdet.SwinTransformer',
    embed_dims=96,
    depths=[2, 2, 6, 2],
    num_heads=[3, 6, 12, 24],
    window_size=7,
    mlp_ratio=4,
    qkv_bias=True,
    qk_scale=None,
    drop_rate=0.0,
    attn_drop_rate=0.0,
    drop_path_rate=0.2,
    patch_norm=True,
    out_indices=[1, 2, 3],
    with_cp=False,
    convert_weights=True,
    init_cfg=dict(
        type='Pretrained',
        checkpoint=  # noqa: E251
        'https://github.com/SwinTransformer/storage/releases/download/v1.0.0/swin_tiny_patch4_window7_224.pth'  # noqa: E501
    )
),

REF_NECK_CFG= dict(
    in_channels=[192, 384, 768],
    out_channels=256,
    start_level=0,
    num_outs=3,
    norm_cfg=dict(type='BN2d', requires_grad=True),
    act_cfg=dict(type='ReLU', inplace=True),
    upsample_cfg=dict(mode='bilinear', align_corners=False)
)

class Swin_OpenPCDet(nn.Module):
    def __init__(self, swin_cfg) -> None:
        super().__init__()
        self.swin = MODELS.build(swin_cfg)
        self.swin.init_weights()
        
    def forward(self, batch_dict):
        x = batch_dict['camera_imgs']
        B, N, C, H, W = x.shape
        x = x.view(B * N, C, H, W).contiguous()
        x = self.swin(x)
        batch_dict["image_features"] = x
        return batch_dict

        
        